# 指定文件编码为UTF-8
# coding: utf-8

"""
hyperparameter_optimization.py 功能说明：
1. 实现神经网络超参数的随机搜索优化
2. 主要优化学习率(lr)和权重衰减系数(weight decay)
3. 使用MNIST数据集进行训练和验证
4. 采用验证集评估模型性能
5. 可视化展示不同超参数组合的训练效果

关键点：
- 使用随机搜索代替网格搜索，提高搜索效率
- 采用对数均匀分布采样超参数
- 使用验证集评估模型泛化能力
- 训练过程封装在Trainer类中
- 可视化展示最佳超参数组合的训练曲线
"""

# 导入系统模块
import sys, os
# 添加父目录到系统路径
sys.path.append(os.pardir)
# 导入NumPy数值计算库
import numpy as np
# 导入matplotlib绘图库
import matplotlib.pyplot as plt
# 从dataset.mnist导入MNIST数据加载函数
from dataset.mnist import load_mnist
# 从common.multi_layer_net导入多层网络类
from common.multi_layer_net import MultiLayerNet
# 从common.util导入数据洗牌函数
from common.util import shuffle_dataset
# 从common.trainer导入训练器类
from common.trainer import Trainer

# 加载MNIST数据集(自动归一化)
(x_train, t_train), (x_test, t_test) = load_mnist(normalize=True)

# 减少训练数据量到500个样本(加速实验)
x_train = x_train[:500]
t_train = t_train[:500]

# 分割验证集(20%的训练数据)
validation_rate = 0.20
validation_num = int(x_train.shape[0] * validation_rate)
# 洗牌数据确保随机性
x_train, t_train = shuffle_dataset(x_train, t_train)
# 划分验证集
x_val = x_train[:validation_num]
t_val = t_train[:validation_num]
# 剩余作为训练集
x_train = x_train[validation_num:]
t_train = t_train[validation_num:]

# 定义训练函数
def __train(lr, weight_decay, epocs=50):
    """
    训练网络并返回准确率列表

    参数:
        lr: 学习率
        weight_decay: 权重衰减系数
        epocs: 训练轮数(默认50)

    返回:
        val_acc_list: 验证集准确率列表
        train_acc_list: 训练集准确率列表
    """
    # 初始化6层神经网络(每层100个神经元)
    network = MultiLayerNet(input_size=784,
                           hidden_size_list=[100, 100, 100, 100, 100, 100],
                           output_size=10,
                           weight_decay_lambda=weight_decay)

    # 创建训练器实例
    trainer = Trainer(network, x_train, t_train, x_val, t_val,
                     epochs=epocs,
                     mini_batch_size=100,
                     optimizer='sgd',
                     optimizer_param={'lr': lr},
                     verbose=False)
    # 开始训练
    trainer.train()

    return trainer.test_acc_list, trainer.train_acc_list

# 超参数随机搜索主循环
optimization_trial = 100  # 随机搜索次数
results_val = {}  # 存储验证集结果
results_train = {}  # 存储训练集结果

for _ in range(optimization_trial):
    # 随机生成超参数(对数尺度)
    weight_decay = 10 ** np.random.uniform(-8, -4)  # 权重衰减系数: 10^-8 ~ 10^-4
    lr = 10 ** np.random.uniform(-6, -2)  # 学习率: 10^-6 ~ 10^-2

    # 训练并获取准确率曲线
    val_acc_list, train_acc_list = __train(lr, weight_decay)

    # 打印当前超参数组合的结果
    print("val acc:" + str(val_acc_list[-1]) + " | lr:" + str(lr) + ", weight decay:" + str(weight_decay))

    # 保存结果(使用超参数组合作为key)
    key = "lr:" + str(lr) + ", weight decay:" + str(weight_decay)
    results_val[key] = val_acc_list
    results_train[key] = train_acc_list

# 可视化最佳超参数组合 =================================================# ... 省略其他代码 ...


# ==== 修改后代码 ====
print("=========== Hyper-Parameter Optimization Result ===========")
graph_draw_num = 20  # 展示前20个最佳结果
col_num = 5  # 每行5个子图
row_num = int(np.ceil(graph_draw_num / col_num))  # 计算需要的行数
i = 0  # 子图索引

# 按验证集准确率排序并绘制结果
for key, val_acc_list in sorted(results_val.items(), key=lambda x:x[1][-1], reverse=True):
    # 打印排名信息
    print("Best-" + str(i+1) + "(val acc:" + str(val_acc_list[-1]) + ") | " + key)

    # 创建子图
    plt.subplot(row_num, col_num, i+1)
    plt.title("Best-" + str(i+1))
    plt.ylim(0.0, 1.0)  # 固定y轴范围

    # 仅在最左侧子图显示y轴刻度
    if i % 5:
        plt.yticks([])
    plt.xticks([])  # 隐藏x轴刻度

    # 绘制验证集和训练集曲线
    x = np.arange(len(val_acc_list))
    plt.plot(x, val_acc_list)  # 验证集(实线)
    plt.plot(x, results_train[key], "--")  # 训练集(虚线)

    i += 1
    # 达到显示数量后停止
    if i >= graph_draw_num:
        break

# 显示图形
plt.show()
